import pdb
import sys

import torch
import torch.nn.functional as F
from torch import nn


class ColorDeformConv2d(nn.Module):
    def __init__(
        self,
        inc,
        outc,
        kernel_size=3,
        padding=1,
        stride=1,
        bias=None,
        modulation=False,
        color_deform=True,
    ):
        """
        Args:
            modulation (bool, optional): If True, Modulated Defomable Convolution (Deformable ConvNets v2).
        """
        super(ColorDeformConv2d, self).__init__()
        self.kernel_size = kernel_size
        self.padding = padding
        self.stride = stride
        self.zero_padding = nn.ZeroPad2d(padding)

        self.conv = nn.Conv2d(
            inc, outc, kernel_size=kernel_size, stride=kernel_size, bias=bias
        )
        self.channel_down = nn.Conv2d(inc * 2, inc, kernel_size=1, stride=1)

        self.p_conv = nn.Conv2d(
            inc,
            2 * kernel_size * kernel_size,
            kernel_size=3,
            padding=1,
            stride=stride,
        )

        self.modulation = modulation
        if modulation:
            self.m_conv = nn.Conv2d(
                inc,
                kernel_size * kernel_size,
                kernel_size=3,
                padding=1,
                stride=stride,
            )

        self.color_deform = color_deform
        if self.color_deform:
            self.c_conv = nn.Conv2d(
                inc,
                kernel_size * kernel_size * inc,
                kernel_size=3,
                padding=1,
                stride=stride,
            )

    def forward(self, x, ref):
        torch.cuda.empty_cache()

        assert (
            x.shape == ref.shape
        ), f"Input shape {x.shape} and reference shape {ref.shape} does not match."

        b, c, h, w = x.size()

        fused = torch.cat([x, ref], dim=1)  # (b, 2c, h, w)
        fused = self.channel_down(fused)  # (b, c, h, w)
        assert (
            fused.shape == x.shape
        ), f"Fused shape {fused.shape} and input shape {x.shape} does not match."

        offset = self.p_conv(fused)  # (b, c, h, w) -> (b, 2*ks*ks, h, w)
        if self.modulation:
            m = torch.sigmoid(
                self.m_conv(fused)
            )  # (b, c, h, w) -> (b, ks*ks, h, w)
        if self.color_deform:
            c_offset = torch.tanh(
                self.c_conv(fused)
            )  # (b, c, h, w)->(b, c*ks*ks, h, w)

        dtype = offset.data.type()
        ks = self.kernel_size  # ks = 3
        N = ks * ks  # N = kernel_size * kernel_size = 9

        if self.padding:  # self.padding = 1
            x = self.zero_padding(
                x
            )  # x.shape: (b, c, h, w) -> (b, c, h+2, w+2)

        # (b, 2N, h, w)
        p = self._get_p(offset, dtype)  # offset.shape: (1, 18, h, w)

        # (b, h, w, 2N)
        p = p.contiguous().permute(0, 2, 3, 1)
        q_lt = p.detach().floor()  # left top
        q_rb = q_lt + 1  # right bottom

        q_lt = torch.cat(
            [
                torch.clamp(q_lt[..., :N], 0, x.size(2) - 1),
                torch.clamp(q_lt[..., N:], 0, x.size(3) - 1),
            ],
            dim=-1,
        ).long()
        q_rb = torch.cat(
            [
                torch.clamp(q_rb[..., :N], 0, x.size(2) - 1),
                torch.clamp(q_rb[..., N:], 0, x.size(3) - 1),
            ],
            dim=-1,
        ).long()
        q_lb = torch.cat([q_lt[..., :N], q_rb[..., N:]], dim=-1)
        q_rt = torch.cat([q_rb[..., :N], q_lt[..., N:]], dim=-1)

        # clip p
        p = torch.cat(
            [
                torch.clamp(p[..., :N], 0, x.size(2) - 1),
                torch.clamp(p[..., N:], 0, x.size(3) - 1),
            ],
            dim=-1,
        )

        # bilinear kernel (b, h, w, N)
        g_lt = (1 + (q_lt[..., :N].type_as(p) - p[..., :N])) * (
            1 + (q_lt[..., N:].type_as(p) - p[..., N:])
        )
        g_rb = (1 - (q_rb[..., :N].type_as(p) - p[..., :N])) * (
            1 - (q_rb[..., N:].type_as(p) - p[..., N:])
        )
        g_lb = (1 + (q_lb[..., :N].type_as(p) - p[..., :N])) * (
            1 - (q_lb[..., N:].type_as(p) - p[..., N:])
        )
        g_rt = (1 - (q_rt[..., :N].type_as(p) - p[..., :N])) * (
            1 + (q_rt[..., N:].type_as(p) - p[..., N:])
        )

        # (b, c, h, w, N)
        x_q_lt = self._get_x_q(x, q_lt, N)
        x_q_rb = self._get_x_q(x, q_rb, N)
        x_q_lb = self._get_x_q(x, q_lb, N)
        x_q_rt = self._get_x_q(x, q_rt, N)

        # (b, c, h, w, ks*ks)
        x_offset = (
            g_lt.unsqueeze(dim=1) * x_q_lt
            + g_rb.unsqueeze(dim=1) * x_q_rb
            + g_lb.unsqueeze(dim=1) * x_q_lb
            + g_rt.unsqueeze(dim=1) * x_q_rt
        )

        if self.color_deform:
            c_offset = c_offset.contiguous().reshape(
                b, c, N, h, w
            )  # (b, c*ks*ks, h, w) -> (b, c, ks*ks, h, w)
            c_offset = c_offset.contiguous().permute(
                0, 1, 3, 4, 2
            )  # (b, c, ks*ks, h, w) -> (b, c, h, w, ks*ks)
            x_offset += c_offset

        # modulation
        if self.modulation:
            m = m.contiguous().permute(0, 2, 3, 1)  # (b, h, w, ks*ks)
            m = m.unsqueeze(dim=1)  # (b, 1, h, w, ks*ks)
            m = torch.cat(
                [m for _ in range(x_offset.size(1))], dim=1
            )  # (b, 1, h, w, ks*ks) -> (b, c, h, w, ks*ks)
            x_offset *= m

        x_offset = self._reshape_x_offset(x_offset, ks)  # (b, c, h*ks, w*ks)
        out = self.conv(x_offset)

        return out

    def _get_p_n(self, N, dtype):
        p_n_x, p_n_y = torch.meshgrid(
            torch.arange(
                -(self.kernel_size - 1) // 2, (self.kernel_size - 1) // 2 + 1
            ),
            torch.arange(
                -(self.kernel_size - 1) // 2, (self.kernel_size - 1) // 2 + 1
            ),
        )
        # (2N, 1)
        p_n = torch.cat([torch.flatten(p_n_x), torch.flatten(p_n_y)], 0)
        p_n = p_n.view(1, 2 * N, 1, 1).type(dtype)

        return p_n

    def _get_p_0(self, h, w, N, dtype):
        p_0_x, p_0_y = torch.meshgrid(
            torch.arange(1, h * self.stride + 1, self.stride),
            torch.arange(1, w * self.stride + 1, self.stride),
        )
        p_0_x = torch.flatten(p_0_x).view(1, 1, h, w).repeat(1, N, 1, 1)
        p_0_y = torch.flatten(p_0_y).view(1, 1, h, w).repeat(1, N, 1, 1)
        p_0 = torch.cat([p_0_x, p_0_y], 1).type(dtype)

        return p_0

    def _get_p(self, offset, dtype):
        N, h, w = offset.size(1) // 2, offset.size(2), offset.size(3)

        # (1, 2N, 1, 1)
        p_n = self._get_p_n(N, dtype)
        # (1, 2N, h, w)
        p_0 = self._get_p_0(h, w, N, dtype)
        p = p_0 + p_n + offset
        return p

    def _get_x_q(self, x, q, N):
        b, h, w, _ = q.size()
        padded_w = x.size(3)
        c = x.size(1)
        # (b, c, h*w)
        x = x.contiguous().view(b, c, -1)

        # (b, h, w, N)
        index = q[..., :N] * padded_w + q[..., N:]  # offset_x*w + offset_y
        # (b, c, h*w*N)
        index = (
            index.contiguous()
            .unsqueeze(dim=1)
            .expand(-1, c, -1, -1, -1)
            .contiguous()
            .view(b, c, -1)
        ).type(torch.int64)

        x_offset = (
            x.gather(dim=-1, index=index).contiguous().view(b, c, h, w, N)
        )

        return x_offset

    @staticmethod
    def _reshape_x_offset(x_offset, ks):
        b, c, h, w, N = x_offset.size()
        x_offset = torch.cat(
            [
                x_offset[..., s : s + ks].contiguous().view(b, c, h, w * ks)
                for s in range(0, N, ks)
            ],
            dim=-1,
        )
        x_offset = x_offset.contiguous().view(b, c, h * ks, w * ks)

        return x_offset
